%% EKF估计非线性离散系统
% W->Q,V->R都给了，只需手动先求出Phi阵和H阵
% 先对于KF，EKF的X一步预测和X状态估值计算不太一样
clear all;
clc;

%% 参数初始化
N = 500;
t = 1:1:N;
X0 = [1; 1];
P0 = [10, 0; 0, 10];
mu = [0, 0];
matQ = [0.01, 0; 0, 0.1];
matR = [1, 0; 0, 0.1];

%% 生成真值和测量值
% 此处用mvnrnd生成噪声阵再加到信息阵里面去
% 之前是直接用mvnrnd生成信息阵
% 不知道有何区别
rng(1);
Wk = mvnrnd(mu, matQ, N)';
Vk = mvnrnd(mu, matR, N)';
realXk = zeros(2, N);
measureXk = zeros(2, N);

realXk(:, 1) = X0;
measureXk(1, 1) = sqrt(realXk(1,1)^2 + realXk(2,1)^2) + Vk(1,1);
measureXk(2, 1) = atan(realXk(1,1) / realXk(2,1)) + Vk(2,1);
for i = 2:1:N
    x1pre = realXk(1, i-1);
    x2pre = realXk(2, i-1);
    realXk(1, i) = x2pre*sin(x1pre) + 0.1*(i-1) + Wk(1, i-1);
    realXk(2, i) = x1pre + (cos(x2pre))^2 - 0.1*(i-1) + Wk(2, i-1);
    measureXk(1, i) = sqrt(realXk(1,i)^2 + realXk(2,i)^2) + Vk(1,i);
    measureXk(2, i) = atan(realXk(1,i) / realXk(2,i)) + Vk(2,i);
end

%% EKF滤波
% 初值
estimateXpre = X0;
meanSquareErrorPre = P0;
estimateX = X0;
% 记录
estimateX_EKF = zeros(2, N);
estimateX_EKF(:, 1) = estimateXpre;
estimateMeanSquareError = zeros(2, N);
estimateMeanSquareError(:, 1) = diag(P0);

for i = 2:1:N
    % 相对于KF，EKF的Phi阵要由状态方程f(X)实时求偏导得到
    matPHIk = [estimateXpre(2)*cos(estimateXpre(1)), sin(estimateXpre(1));
               1, -sin(2*estimateXpre(2))];
    % 相对于KF，EKF的一部预测方程使用非线性的状态方程f(X)
    estimateXpre = [estimateXpre(2)*sin(estimateXpre(1)) + 0.1*(i-1);
                    estimateXpre(1) + (cos(estimateXpre(2)))^2 - 0.1*(i-1)];
    % 一步预测Pk
    meanSquareErrorPre = matPHIk*meanSquareErrorPre*matPHIk' + matQ;
    % 相对于KF，EKF的Hk阵要由观测方程h(X)实时求偏导得到
    tmpX_2 = estimateXpre(1)^2 + estimateXpre(2)^2;
    matHk = [estimateXpre(1)/sqrt(tmpX_2), estimateXpre(2)/sqrt(tmpX_2);
             estimateXpre(2)/tmpX_2, -estimateXpre(1)/tmpX_2];
    matHkT = matHk';
    % 增益K
    gainK = meanSquareErrorPre*matHkT/(matHk*meanSquareErrorPre*matHkT + matR);
    % X估计，相对于KF，此处使用观测方程h(X)
    hx = [sqrt(estimateXpre(1)^2 + estimateXpre(2)^2);
          atan(estimateXpre(1) / estimateXpre(2))];
    estimateX = estimateXpre + gainK*(measureXk(:,i) - hx);
    % P估计
    tmp = eye(2) - gainK*matHk;
    meanSquareErrorK = tmp*meanSquareErrorPre*tmp' + gainK*matR*gainK';
    
    estimateXpre = estimateX;
    meanSquareErrorPre = meanSquareErrorK;
    % 记录
    estimateX_EKF(:, i) = estimateX;
    estimateMeanSquareError(:, i) = diag(meanSquareErrorK);
end

%% 画图
% 真值
figure;
plot(t, realXk);
legend('x1','x2');
title('真值');
% 测量值
figure;
plot(t, measureXk);
legend('x1','x2');
title('测量值');
% 估计值
figure;
plot(t, estimateX_EKF);
legend('x1','x2');
title('EKF估计值');
% 估计误差值
figure;
plot(t, estimateX_EKF - realXk);
legend('x1','x2');
title('估计误差值');
% 真值vs估计值
figure;
plot(t, realXk);
hold on;
plot(t, estimateX_EKF);
legend('realx1','realx2','EKFx1','EKFx2');
title('真值vs估计值');
% 均方差
figure;
plot(t, estimateMeanSquareError);
legend('x1','x2');
title('均方差');
